import os

from langchain_core.prompts import PromptTemplate
from langchain_ollama import OllamaLLM

from Knowledge.pg_store import insert_into_pg
from Knowledge.md_mineru_parse import parse_markdown
from Knowledge.milvus_store import insert_into_milvus
from Search.route import select_knowledge_base

llm = OllamaLLM(model="deepseek-r1:70b", base_url="http://192.168.7.3:11434")


def main():
    md_path = r"C:\Desktop\各类知识\知识手册\2019-03-08_中科宇图科技股份有限公司精准治霾基础知识学习手册_1.0.0.md"
    filename = os.path.basename(md_path)
    if not os.path.exists(md_path):
        print("Markdown 文件不存在，请检查路径")
        return

    with open(md_path, "r", encoding="utf-8") as f:
        md_text = f.read()

    # 解析 Markdown
    parent_docs, child_docs = parse_markdown(md_text, filename)

    # 存入 PostgreSQL（父文档）
    insert_into_pg(parent_docs)

    # 存入 Milvus（子文档）
    insert_into_milvus(child_docs)

    print("✅ 数据存储完成！可以开始检索。")


if __name__ == "__main__":
    # main()
    query = "最新的城市环境空气质量排名技术规定有哪些？"
    results = select_knowledge_base(query)
    print("\n🔍 查询结果：")
    print(results)




    # 将多个文档合并为一个文本
    if isinstance(results, list):
        text = "\n\n".join(results)
    elif isinstance(results, str):
        text = results
    else:
        text = str(results)
    prompt = PromptTemplate(
        input_variables=["query", "text"],
        template=(
            "你是一名专业的环境科学专家。\n\n"
            "请根据以下提供的文本，回答用户的问题。\n"
            "如果文本中包含相关信息，请使用专业、准确的语言回答。\n"
            "如果文本中没有足够的信息，请直接回答：我不知道。\n\n"
            "### 用户问题：\n{query}\n\n"
            "### 参考文本：\n{text}\n\n"
            "请输出清晰、专业的回答，不要输出无关信息。"
        ),
    )
    response = llm.invoke(prompt.format(query=query, text=text)).strip()
    print("\n🤖 LLM 回答：")
    print(response)
